#coding=utf-8
#! /bin/env python

from dffpy import *

FILENAME = r"CPMDTRAJEC.xyz"
PBC_BOX = "PBC: 15, 20, 13, 90, 90, 90"

last_step = 0

def formula(mol):
    counter = {}
    nAtoms = mol.count_atoms()
    for i in range(nAtoms):
        atom = mol.get_atom(i)
        if counter.has_key(atom.symbol):
            counter[atom.symbol] += 1
        else:
            counter[atom.symbol] = 1
    result = ""
    if counter.has_key("C"):
        result += "C"
        if counter["C"] > 1:
            result += str(counter["C"])
        del counter["C"]
    if counter.has_key("H"):
        result += "H"
        if counter["H"] > 1:
            result += str(counter["H"])
        del counter["H"]
    for (k, v) in counter.iteritems():
        result += k
        if v > 1:
            result += str(v)
    return result
    

def count_bonds(atoms, bonds):
    print "Has", len(bonds), "bonds"
    counter = {}
    for i in bonds:
        sym1 = atoms[i.atom1].symbol
        sym2 = atoms[i.atom2].symbol
        if sym1 > sym2:
            (sym1, sym2) = (sym2, sym1)
        symbol = sym1 + "--" + sym2
        if counter.has_key(symbol):
            counter[symbol] += 1
        else:
            counter[symbol] = 1
    print counter

FRAGMENT_COUNT = {}
def count_mols(atoms, bonds, step):
    global FRAGMENT_COUNT
    mol = MolecularModel.setup(atoms, bonds)
    mols = get_connected_parts(mol)
    print "Has", len(mols), "mols"
    counter = {}
    for i in mols:
        form = formula(i)
        if counter.has_key(form):
            counter[form] += 1
        else:
            counter[form] = 1
    print counter

    for (k, v) in counter.iteritems():
        if not FRAGMENT_COUNT.has_key(k):
            FRAGMENT_COUNT[k] = {}
        FRAGMENT_COUNT[k][step] = v


def parse_block(block, pbc):
    global last_step
    if block[0].startswith(" STEP:"):
        step = int(block[0][8:].strip())
        if step != last_step:
            last_step = step
            print " ========== Step %8d ========== " % step 
            atoms = []
            for i in block[1:]:
                tokens = i.split()
                if len(tokens) == 4:
                    a = Atom(tokens[0].strip())
                    a.coord.x = float(tokens[1])
                    a.coord.y = float(tokens[2])
                    a.coord.z = float(tokens[3])
                    atoms.append(a)
            bonds = guess_bonds(atoms, pbc, 0.1)
            if len(bonds):
                count_bonds(atoms, bonds)
                count_mols(atoms, bonds, step)
            

if __name__ == '__main__':
    pbc = PBC.parse(PBC_BOX)
    with open(FILENAME, "rt") as file:
        block = []
        for line in file:
            if line.startswith(" STEP:"):
                parse_block(block, pbc)
                block = [line]
            else:
                block.append(line)
        parse_block(block, pbc)

    print
    print
    print
    for (k, v) in FRAGMENT_COUNT.iteritems():
        print " ========== %10s ========== " % k
        counter[k] = 0
        for (kk, vv) in v.iteritems():
            print kk, vv
        print

    print
    print
    print
    for (k, v) in FRAGMENT_COUNT.iteritems():
        print k, sum(v.itervalues())
        
